'''
Author: Yan Jianhang
Date: 2022-04-12 18:36:43
LastEditors: Yan Jianhang
LastEditTime: 2022-04-12 18:43:33
FilePath: /STransGrasp/lib/tf/utils.py
Description: 
'''
import sys
import open3d as o3d
import torch

def batch_pairwise_dist(x, y, use_cuda=False):
    bs, num_points_x, points_dim = x.size()
    _, num_points_y, _ = y.size()
    xx = torch.bmm(x, x.transpose(2, 1))
    yy = torch.bmm(y, y.transpose(2, 1))
    zz = torch.bmm(x, y.transpose(2, 1))
    if use_cuda:
        dtype = torch.cuda.LongTensor
    else:
        dtype = torch.LongTensor
    diag_ind_x = torch.arange(0, num_points_x).type(dtype)
    diag_ind_y = torch.arange(0, num_points_y).type(dtype)
    # brk()
    rx = xx[:, diag_ind_x, diag_ind_x].unsqueeze(1).expand_as(zz.transpose(2, 1))
    ry = yy[:, diag_ind_y, diag_ind_y].unsqueeze(1).expand_as(zz)
    P = (rx.transpose(2, 1) + ry - 2 * zz)
    return P

def get_grasp_hand_lineset(grasp_points):
    grasp_hand = o3d.geometry.LineSet()
    grasp_hand.points = o3d.utility.Vector3dVector(grasp_points)
    #     Food:        (0,  1),  2,     3  --            (FF4, FF3),   FF2,     FF1
    #     Middle:      (4,  5),  6,     7  --            (MF4, MF3),   MF2,     MF1
    #     Ring:        (8,  9), 10,    11  --            (RF4, RF3),   RF2,     RF1
    #     little: 12,  (13,14), 15,    16  --    LF5,    (LF4, LF3),   LF2,     LF1
    #     Thumb:       (17,18),(19,20),21  --            (TF5, TF4),(TF3, TF2), TF1

    lines = [[0, 1], [2, 3], [3, 4], [4, 23],
            [0, 5], [6, 7], [7, 8], [8, 24],
            [0, 9], [10, 11], [11, 12], [12, 25],
            [0, 13], [13, 15], [15, 16], [16, 17], [17, 26],
            [0, 18], [19, 21], [21, 22], [22, 27]]
    grasp_hand.lines = o3d.utility.Vector2iVector(lines)

    return grasp_hand